// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/row/unsafe_row_fast.h>

namespace kumo::pollux::row {
    namespace {
        static const int32_t kFieldWidth = 8;
        using TRowSize = uint32_t;

        int32_t alignBits(int32_t numBits) {
            return bits::nwords(numBits) * 8;
        }

        int32_t alignBytes(int32_t numBytes) {
            return bits::roundUp(numBytes, 8);
        }

        bool isFixedWidth(const TypePtr &type) {
            return type->isFixedWidth() && !type->isLongDecimal();
        }
    } // namespace

    // static
    std::optional<int32_t> UnsafeRowFast::fixedRowSize(const RowTypePtr &rowType) {
        for (const auto &child: rowType->children()) {
            if (!isFixedWidth(child)) {
                return std::nullopt;
            }
        }

        const size_t numFields = rowType->size();
        const size_t nullLength = alignBits(numFields);

        return nullLength + numFields * kFieldWidth;
    }

    UnsafeRowFast::UnsafeRowFast(const RowVectorPtr &vector)
        : typeKind_{vector->type_kind()}, decoded_{*vector} {
        initialize(vector->type());
    }

    UnsafeRowFast::UnsafeRowFast(const VectorPtr &vector)
        : typeKind_{vector->type_kind()}, decoded_{*vector} {
        initialize(vector->type());
    }

    void UnsafeRowFast::initialize(const TypePtr &type) {
        auto base = decoded_.base();
        switch (typeKind_) {
            case TypeKind::ARRAY: {
                auto arrayBase = base->as<ArrayVector>();
                children_.push_back(UnsafeRowFast(arrayBase->elements()));
                childIsFixedWidth_.push_back(isFixedWidth(arrayBase->elements()->type()));
                break;
            }
            case TypeKind::MAP: {
                auto mapBase = base->as<MapVector>();
                children_.push_back(UnsafeRowFast(mapBase->mapKeys()));
                children_.push_back(UnsafeRowFast(mapBase->mapValues()));
                childIsFixedWidth_.push_back(isFixedWidth(mapBase->mapKeys()->type()));
                childIsFixedWidth_.push_back(isFixedWidth(mapBase->mapValues()->type()));
                break;
            }
            case TypeKind::ROW: {
                auto rowBase = base->as<RowVector>();
                for (const auto &child: rowBase->children()) {
                    children_.push_back(UnsafeRowFast(child));
                    childIsFixedWidth_.push_back(isFixedWidth(child->type()));
                }

                rowNullBytes_ = alignBits(type->size());
                break;
            }
            case TypeKind::BOOLEAN:
                valueBytes_ = 1;
                fixedWidthTypeKind_ = true;
                break;
            case TypeKind::TINYINT:
                [[fallthrough]];
            case TypeKind::SMALLINT:
                [[fallthrough]];
            case TypeKind::INTEGER:
                [[fallthrough]];
            case TypeKind::BIGINT:
                [[fallthrough]];
            case TypeKind::REAL:
                [[fallthrough]];
            case TypeKind::DOUBLE:
                [[fallthrough]];
            case TypeKind::UNKNOWN:
                valueBytes_ = type->cppSizeInBytes();
                fixedWidthTypeKind_ = true;
                supportsBulkCopy_ = decoded_.isIdentityMapping();
                break;
            case TypeKind::TIMESTAMP:
                valueBytes_ = sizeof(int64_t);
                fixedWidthTypeKind_ = true;
                break;
            case TypeKind::HUGEINT:
                [[fallthrough]];
            case TypeKind::VARCHAR:
                [[fallthrough]];
            case TypeKind::VARBINARY:
                // Nothing to do.
                break;
            default:
                POLLUX_UNSUPPORTED("Unsupported type: {}", type->toString());
        }
    }

    void UnsafeRowFast::serializedRowSizes(
        const melon::Range<const vector_size_t *> &rows,
        vector_size_t **sizes) const {
        if (const auto fixedRowSize =
                UnsafeRowFast::fixedRowSize(as_row_type(decoded_.base()->type()))) {
            for (const auto row: rows) {
                *sizes[row] = fixedRowSize.value() + sizeof(TRowSize);
            }
            return;
        }

        for (const auto &row: rows) {
            *sizes[row] = rowSize(row) + sizeof(TRowSize);
        }
    }

    int32_t UnsafeRowFast::rowSize(vector_size_t index) const {
        return rowRowSize(index);
    }

    int32_t UnsafeRowFast::variableWidthRowSize(vector_size_t index) const {
        switch (typeKind_) {
            case TypeKind::VARCHAR:
                [[fallthrough]];
            case TypeKind::VARBINARY: {
                auto value = decoded_.value_at<StringView>(index);
                return alignBytes(value.size());
            }
            case TypeKind::HUGEINT:
                return DecimalUtil::getByteArrayLength(decoded_.value_at<int128_t>(index));
            case TypeKind::ARRAY:
                return arrayRowSize(index);
            case TypeKind::MAP:
                return mapRowSize(index);
            case TypeKind::ROW:
                return rowRowSize(index);
            default:
                POLLUX_UNREACHABLE(
                    "Unexpected type kind: {}", mapTypeKindToName(typeKind_));
        };
    }

    bool UnsafeRowFast::is_null_at(vector_size_t index) const {
        return decoded_.is_null_at(index);
    }

    int32_t UnsafeRowFast::serialize(vector_size_t index, char *buffer) const {
        return serializeRow(index, buffer);
    }

    void UnsafeRowFast::serializeFixedWidth(vector_size_t index, char *buffer)
    const {
        POLLUX_DCHECK(fixedWidthTypeKind_);
        switch (typeKind_) {
            case TypeKind::BOOLEAN:
                *reinterpret_cast<bool *>(buffer) = decoded_.value_at<bool>(index);
                break;
            case TypeKind::TIMESTAMP:
                *reinterpret_cast<int64_t *>(buffer) =
                        decoded_.value_at<Timestamp>(index).toMicros();
                break;
            default:
                memcpy(
                    buffer,
                    decoded_.data<char>() + decoded_.index(index) * valueBytes_,
                    valueBytes_);
        }
    }

    void UnsafeRowFast::serializeFixedWidth(
        vector_size_t offset,
        vector_size_t size,
        char *buffer) const {
        POLLUX_DCHECK(supportsBulkCopy_);
        // decoded_.data<char>() can be null if all values are null.
        if (decoded_.data<char>()) {
            memcpy(
                buffer,
                decoded_.data<char>() + decoded_.index(offset) * valueBytes_,
                valueBytes_ * size);
        }
    }

    int32_t UnsafeRowFast::serializeVariableWidth(vector_size_t index, char *buffer)
    const {
        switch (typeKind_) {
            case TypeKind::VARCHAR:
                [[fallthrough]];
            case TypeKind::VARBINARY: {
                auto value = decoded_.value_at<StringView>(index);
                memcpy(buffer, value.data(), value.size());
                return value.size();
            }
            case TypeKind::HUGEINT: {
                auto value = decoded_.value_at<int128_t>(index);
                return DecimalUtil::toByteArray(value, buffer);
            }
            case TypeKind::ARRAY:
                return serializeArray(index, buffer);
            case TypeKind::MAP:
                return serializeMap(index, buffer);
            case TypeKind::ROW:
                return serializeRow(index, buffer);
            default:
                POLLUX_UNREACHABLE(
                    "Unexpected type kind: {}", mapTypeKindToName(typeKind_));
        };
    }

    int32_t UnsafeRowFast::arrayRowSize(vector_size_t index) const {
        auto baseIndex = decoded_.index(index);

        // array size | null bits | fixed-width data | variable-width data
        auto arrayBase = decoded_.base()->as_unchecked<ArrayVector>();
        auto offset = arrayBase->offsetAt(baseIndex);
        auto size = arrayBase->sizeAt(baseIndex);

        return arrayRowSize(children_[0], offset, size, childIsFixedWidth_[0]);
    }

    int32_t UnsafeRowFast::serializeArray(vector_size_t index, char *buffer) const {
        auto baseIndex = decoded_.index(index);

        // array size | null bits | fixed-width data | variable-width data
        auto arrayBase = decoded_.base()->as_unchecked<ArrayVector>();
        auto offset = arrayBase->offsetAt(baseIndex);
        auto size = arrayBase->sizeAt(baseIndex);

        return serializeAsArray(
            children_[0], offset, size, childIsFixedWidth_[0], buffer);
    }

    int32_t UnsafeRowFast::mapRowSize(vector_size_t index) const {
        auto baseIndex = decoded_.index(index);

        //  size of serialized keys array in bytes | <keys array> | <values array>

        auto mapBase = decoded_.base()->as_unchecked<MapVector>();
        auto offset = mapBase->offsetAt(baseIndex);
        auto size = mapBase->sizeAt(baseIndex);

        return kFieldWidth +
               arrayRowSize(children_[0], offset, size, childIsFixedWidth_[0]) +
               arrayRowSize(children_[1], offset, size, childIsFixedWidth_[1]);
    }

    int32_t UnsafeRowFast::serializeMap(vector_size_t index, char *buffer) const {
        auto baseIndex = decoded_.index(index);

        //  size of serialized keys array in bytes | <keys array> | <values array>

        auto mapBase = decoded_.base()->as_unchecked<MapVector>();
        auto offset = mapBase->offsetAt(baseIndex);
        auto size = mapBase->sizeAt(baseIndex);

        int32_t serializedBytes = kFieldWidth;

        auto keysSerializedBytes = serializeAsArray(
            children_[0],
            offset,
            size,
            childIsFixedWidth_[0],
            buffer + serializedBytes);
        serializedBytes += keysSerializedBytes;

        auto valuesSerializedBytes = serializeAsArray(
            children_[1],
            offset,
            size,
            childIsFixedWidth_[1],
            buffer + serializedBytes);
        serializedBytes += valuesSerializedBytes;

        // Write the size of serialized keys.
        *reinterpret_cast<int64_t *>(buffer) = keysSerializedBytes;

        return serializedBytes;
    }

    int32_t UnsafeRowFast::arrayRowSize(
        const UnsafeRowFast &elements,
        vector_size_t offset,
        vector_size_t size,
        bool fixedWidth) const {
        int32_t nullBytes = alignBits(size);

        int32_t rowSize = kFieldWidth + nullBytes;
        if (fixedWidth) {
            return rowSize + size * elements.valueBytes();
        }

        rowSize += size * kFieldWidth;

        for (auto i = 0; i < size; ++i) {
            if (!elements.is_null_at(offset + i)) {
                rowSize += alignBytes(elements.variableWidthRowSize(offset + i));
            }
        }

        return rowSize;
    }

    int32_t UnsafeRowFast::serializeAsArray(
        const UnsafeRowFast &elements,
        vector_size_t offset,
        vector_size_t size,
        bool fixedWidth,
        char *buffer) const {
        // array size | null bits | fixed-width data | variable-width data

        // Write array size.
        *reinterpret_cast<int64_t *>(buffer) = size;

        int32_t nullBytes = alignBits(size);

        int32_t nullsOffset = sizeof(int64_t);
        int32_t fixedWidthOffset = nullsOffset + nullBytes;

        auto childSize = fixedWidth ? elements.valueBytes() : kFieldWidth;

        int64_t variableWidthOffset = fixedWidthOffset + size * childSize;

        if (elements.supportsBulkCopy_) {
            if (elements.decoded_.may_have_nulls()) {
                for (auto i = 0; i < size; ++i) {
                    if (elements.is_null_at(offset + i)) {
                        bits::setBit(buffer + nullsOffset, i, true);
                    }
                }
            }
            elements.serializeFixedWidth(offset, size, buffer + fixedWidthOffset);
            return variableWidthOffset;
        }

        for (auto i = 0; i < size; ++i) {
            if (elements.is_null_at(offset + i)) {
                bits::setBit(buffer + nullsOffset, i, true);
            } else {
                if (fixedWidth) {
                    elements.serializeFixedWidth(
                        offset + i, buffer + fixedWidthOffset + i * childSize);
                } else {
                    auto serializedBytes = elements.serializeVariableWidth(
                        offset + i, buffer + variableWidthOffset);

                    // Write size and offset.
                    uint64_t sizeAndOffset = variableWidthOffset << 32 | serializedBytes;
                    reinterpret_cast<uint64_t *>(buffer + fixedWidthOffset)[i] =
                            sizeAndOffset;

                    variableWidthOffset += alignBytes(serializedBytes);
                }
            }
        }
        return variableWidthOffset;
    }

    int32_t UnsafeRowFast::rowRowSize(vector_size_t index) const {
        auto childIndex = decoded_.index(index);

        const auto numFields = children_.size();
        int32_t size = rowNullBytes_ + numFields * kFieldWidth;

        for (auto i = 0; i < numFields; ++i) {
            if (!childIsFixedWidth_[i] && !children_[i].is_null_at(childIndex)) {
                size += alignBytes(children_[i].variableWidthRowSize(childIndex));
            }
        }

        return size;
    }

    int32_t UnsafeRowFast::serializeRow(vector_size_t index, char *buffer) const {
        auto childIndex = decoded_.index(index);

        int64_t variableWidthOffset = rowNullBytes_ + kFieldWidth * children_.size();

        for (auto i = 0; i < children_.size(); ++i) {
            auto &child = children_[i];

            // Write null bit.
            if (child.is_null_at(childIndex)) {
                bits::setBit(buffer, i, true);
                continue;
            }

            // Write value.
            if (childIsFixedWidth_[i]) {
                child.serializeFixedWidth(
                    childIndex, buffer + rowNullBytes_ + i * kFieldWidth);
            } else {
                auto size = child.serializeVariableWidth(
                    childIndex, buffer + variableWidthOffset);
                // Write size and offset.
                uint64_t sizeAndOffset = variableWidthOffset << 32 | size;
                reinterpret_cast<uint64_t *>(buffer + rowNullBytes_)[i] = sizeAndOffset;

                variableWidthOffset += alignBytes(size);
            }
        }

        return variableWidthOffset;
    }
} // namespace kumo::pollux::row
